package org.apache.commons.graph.algorithms.spanning

import org.apache.commons.graph.Graph
import org.apache.commons.graph.Mapper
import org.apache.commons.graph.SpanningTree
import org.apache.commons.graph.collections.DisjointSet
import org.apache.commons.graph.collections.FibonacciHeap
import org.apache.commons.graph.model.MutableSpanningTree
import org.apache.commons.graph.utils.Assertions
import org.apache.commons.graph.weight.OrderedMonoid

/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements.  See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership.  The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License.  You may obtain a copy of the License at
*
*   http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied.  See the License for the
* specific language governing permissions and limitations
* under the License.
*/

/**
 * Spanning Tree algoritms selector.
 *
 * @param <V> the Graph vertices type
 * @param <W> the weight type
 * @param <WE> the Graph weighted edges type
</WE></W></V> */
interface SpanningTreeAlgorithmSelector<V, W, WE> {
    /**
     * Applies the [Boruvka](https://en.wikipedia.org/wiki/Bor%C5%AFvka's_algorithm)'s algorithm.
     *
     * @param <WO> the type of weight operations
     * @param weightOperations the class responsible for operations on weights
     * @return the calculated spanning tree
    </WO> */
    fun <WO : OrderedMonoid<W>> applyingBoruvkaAlgorithm(weightOperations: WO): SpanningTree<V, WE, W>

    /**
     * Applies the [Kruskal](https://en.wikipedia.org/wiki/Kruskal%27s_algorithm)'s algorithm.
     *
     * @param <WO> the type of weight operations
     * @param weightOperations the class responsible for operations on weights
     * @return the calculated spanning tree
    </WO> */
    fun <WO : OrderedMonoid<W>> applyingKruskalAlgorithm(weightOperations: WO): SpanningTree<V, WE, W>

    /**
     * Applies the [Prim](https://en.wikipedia.org/wiki/Prim%27s_algorithm)'s algorithm.
     *
     * @param <WO> the type of weight operations
     * @param weightOperations the class responsible for operations on weights
     * @return the calculated spanning tree
    </WO> */
    fun <WO : OrderedMonoid<W>> applyingPrimAlgorithm(weightOperations: WO): SpanningTree<V, WE, W>
}


/**
 * [SpanningTreeAlgorithmSelector] implementation.
 *
 * @param <V> the Graph vertices type
 * @param <W> the weight type
 * @param <WE> the Graph weighted edges type
</WE></W></V> */
internal class DefaultSpanningTreeAlgorithmSelector<V, W, WE>
    (
    /** The graph.  */
    private val graph: Graph<V, WE>,
    private val weightedEdges: Mapper<WE, W>,
    /** The start vertex.  */
    private val source: V
) : SpanningTreeAlgorithmSelector<V, W, WE> {

    /** {@inheritDoc}  */
    override fun <WO : OrderedMonoid<W>> applyingBoruvkaAlgorithm(weightOperations: WO): SpanningTree<V, WE, W> {
        /*
         * <pre>
         * procedure Boruvka MST(G(V; E)):
         *     T <= V
         *     while |T| < n do
         *         for all connected component C in T do
         *             e <= the smallest-weight edge from C to another component in T
         *             if e not exists in T then
         *                 T <= T u {e}
         *             end if
         *         end for
         *     end while
         * <pre>
         */

//        Assertions.checkNotNull(
//            weightOperations,
//            "The Boruvka algorithm cannot be calculated with null weight operations"
//        )

        val spanningTree = MutableSpanningTree<V, WE, W>(weightOperations, weightedEdges)

        val components: MutableSet<SuperVertex<V, W, WE>> = HashSet(graph.order)

        val mapping: MutableMap<V, SuperVertex<V, W, WE>> = HashMap(graph.order)

        for (v in graph.vertices) {
            // create a super vertex for each vertex
            val sv = SuperVertex<V, W, WE>(v, graph, WeightedEdgesComparator(weightOperations, weightedEdges))

            components.add(sv)

            // add a mapping for each vertex to its corresponding super vertex
            mapping[v] = sv

            // add each vertex to the spanning tree
            spanningTree.addVertex(v)
        }

        while (components.size > 1) {
            val edges: MutableList<WE> = mutableListOf()
            for (sv in components) {
                // get the minimum edge for each component to any other component
                val edge = sv.minimumWeightEdge
                if (edge != null) {
                    edges.add(edge)
                }
            }

            // if there is no edge anymore for a component, and there is still more than 1 component,
            // the graph is unconnected
            Assertions.checkState(edges.isNotEmpty() || components.size == 1, "unconnected graph")

            for (edge in edges) {
                val pair = graph.getVertices(edge)
                val head = pair!!.head
                val tail = pair.tail

                // find the super vertices corresponding to this edge
                val headSv = mapping[head]!!
                val tailSv = mapping[tail]!!

                // merge them, if they are not the same
                if (headSv != tailSv) {
                    headSv.merge(tailSv)

                    // update the mapping for each merged vertex
                    for (v in tailSv) {
                        mapping[v] = headSv
                    }

                    // remove the merged super vertex from the components set
                    components.remove(tailSv)

                    // add the edge to the spanning tree
                    if (spanningTree.getVertices(edge) == null) {
                        spanningTree.addEdge(head, edge, tail)
                    }
                }
            }
        }

        return spanningTree
    }

    /**
     * {@inheritDoc}
     */
    override fun <WO : OrderedMonoid<W>> applyingKruskalAlgorithm(weightOperations: WO): SpanningTree<V, WE, W> {
        Assertions.checkNotNull(
            weightOperations,
            "The Kruskal algorithm cannot be calculated with null weight operations"
        )
        val settledNodes: MutableSet<V> = HashSet()

        val orderedEdges = FibonacciHeap(WeightedEdgesComparator(weightOperations, weightedEdges))

        for (edge in graph.edges) {
            orderedEdges.add(edge)
        }

        val disjointSet = DisjointSet<V>()

        val spanningTree = MutableSpanningTree<V, WE, W>(weightOperations, weightedEdges)

        // fill the spanning tree with vertices.
        for (v in graph.vertices) {
            spanningTree.addVertex(v)
        }

        while (!orderedEdges.isEmpty() && settledNodes.size < graph.order) {
            val edge = orderedEdges.removeAndGet()!!

            val vertices = graph.getVertices(edge)
            val head = vertices!!.head
            val tail = vertices.tail
            settledNodes.add(head)
            settledNodes.add(tail)

            if (disjointSet.find(head) != disjointSet.find(tail)) {
                disjointSet.union(head, tail)
                spanningTree.addEdge(head, edge, tail)
            }
        }

        return spanningTree
    }

    /**
     * {@inheritDoc}
     */
    override fun <WO : OrderedMonoid<W>> applyingPrimAlgorithm(weightOperations: WO): SpanningTree<V, WE, W> {
        Assertions.checkNotNull(weightOperations, "The Prim algorithm cannot be calculated with null weight operations")

        val shortestEdges = org.apache.commons.graph.algorithms.spanning.ShortestEdges(
            graph,
            source,
            weightOperations,
            weightedEdges
        )

        val unsettledNodes = FibonacciHeap(shortestEdges)
        unsettledNodes.add(source)

        val settledEdges: MutableSet<WE?> = HashSet()

        // extract the node with the shortest distance
        while (!unsettledNodes.isEmpty()) {
            val vertex = unsettledNodes.removeAndGet()!!

            for (v in graph.getConnectedVertices(vertex)) {
                val edge = graph.getEdge(vertex, v)!!
                // if the edge has not been already visited and its weight is
                // less then the current Vertex weight
                val weightLessThanCurrent =
                    !shortestEdges.hasWeight(v)
                            || weightOperations.compare(weightedEdges.map(edge), shortestEdges.getWeight(v)) < 0
                if (settledEdges.add(edge) && weightLessThanCurrent) {
                    if (!unsettledNodes.contains(v)) {
                        unsettledNodes.add(v)
                    }

                    shortestEdges.addPredecessor(v, edge)
                }
            }
        }

        return shortestEdges.createSpanningTree()
    }
}
